import numpy as np
import netCDF4 as nc
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from matplotlib.colors import ListedColormap

# 云类型和颜色映射
cloud_types = {
    2: "Water Type",
    3: "Super Cooled Type",
    4: "Mixed Type",
    5: "Ice Type",
    6: "Cirrus Type",
    7: "Overlap Type"
}

# 配置颜色
colors = {
    2: "#F9F8CA",
    3: "#96D2B0",
    4: "#35B9C5",
    5: "#2681B6",
    6: "#1E469B",
    7: "#080f40"
}

# 加载数据
nf = nc.Dataset(r'/mnt/datastore/liudddata/result/20200104new/2020010104_predicted_2d.nc', 'r')
clt_data = np.ma.getdata(nf.variables['clt'][:])  # 读取云类型
lat = nf.variables['lat'][:]
lon = nf.variables['lon'][:]
nf.close()

# 处理无效值（假设填充值为-999或其他特殊值）
clt_data = np.ma.masked_invalid(clt_data)  # 自动屏蔽NaN和inf
clt_data = np.clip(clt_data, 2, 7)  # 确保值在2到7之间

# 处理 lon 和 lat 中的无效值
lon = np.ma.getdata(lon)
lat = np.ma.getdata(lat)
lon = np.nan_to_num(lon, nan=np.nanmean(lon))
lat = np.nan_to_num(lat, nan=np.nanmean(lat))

# 创建地图画布，使用正交投影
fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(1, 1, 1, projection=ccrs.Orthographic(central_latitude=0, central_longitude=104.7))

# 添加地理特征
ax.add_feature(cfeature.COASTLINE)
ax.add_feature(cfeature.BORDERS, linestyle=':')
ax.add_feature(cfeature.LAND, facecolor='white')

# 添加经纬度线
parallels = range(-90, 91, 30)  # 每30度一个纬线
meridians = range(-180, 181, 30)  # 每30度一个经线
gl = ax.gridlines(draw_labels=True, color='gray', linestyle='--', xlocs=meridians, ylocs=parallels)
gl.top_labels = False
gl.right_labels = False
gl.xlabel_style = {'size': 10}
gl.ylabel_style = {'size': 10}

# 创建自定义颜色映射
cmap_colors = [colors[i] for i in sorted(colors.keys())]
cmap = ListedColormap(cmap_colors)

# 绘制云类型数据
im = ax.pcolormesh(lon, lat, clt_data, transform=ccrs.PlateCarree(),
                   cmap=cmap, vmin=2, vmax=7)

# 添加颜色条
cbar_ax = fig.add_axes([0.88, 0.15, 0.03, 0.7])  # 调整颜色条的位置和尺寸
cbar = plt.colorbar(im, cax=cbar_ax, orientation='vertical')
cbar.set_ticks(np.arange(2, 8))  # 设置刻度
cbar.set_ticklabels([cloud_types[i] for i in sorted(cloud_types.keys())])  # 设置刻度标签
cbar.set_label('Cloud Type', fontsize=12)

# 显示图形
plt.show()